from langchain_community.agent_toolkits import load_tools
from langchain.agents import AgentExecutor
from langchain.agents.output_parsers import ReActSingleInputOutputParser
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.tools.render import render_text_description
from langchain import hub
from customize.get_ollama import GetOllama

llm = GetOllama(ip=GetOllama.ailab_linux_ip, model_name="qwen2.5:14b", model_type=1)()
llm_with_stop = llm.bind(stop=['\nObservation'])

tools = load_tools.load_tools(["llm-math"], llm=llm) #"ddg-search",

prompt = hub.pull("hwchase17/react")
prompt = prompt.partial(
    tools=render_text_description(tools),
    tool_names=", ".join([t.name for t in tools])
)
agent = (
    {
        "input": lambda x: x['input'],
        "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
    }
    | prompt
    | llm_with_stop
    | ReActSingleInputOutputParser()
)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
agent_executor.invoke({"input": "从溥仪退位到新中国成立经历了多少年？"})
